# betmr.py

import random
from config import THRESHOLD_SET, SAMPLE_RATIO

class BETMRExperienceReplay:
    def __init__(self):
        self.normal_pool = []     # 正常经验池
        self.expert_pool = []     # 专家轨迹池
        self.threshold = THRESHOLD_SET[0]  # 初始阈值

    def add_experience(self, experience, total_reward):
        # 添加经验到正常池
        self.normal_pool.append(experience)
        # 判断是否为专家轨迹
        if total_reward > self.threshold:
            self.expert_pool.append(experience)
            self._update_threshold(total_reward)

    def _update_threshold(self, reward):
        # 动态更新阈值
        for i in range(len(THRESHOLD_SET) - 1):
            if THRESHOLD_SET[i] < reward < THRESHOLD_SET[i+1]:
                self.threshold = THRESHOLD_SET[i+1]
                break

    def sample(self, batch_size):
        # 按比例采样
        sample_size = int(batch_size * SAMPLE_RATIO)
        normal_samples = random.sample(self.normal_pool, sample_size)
        expert_samples = random.sample(self.expert_pool, batch_size - sample_size)
        return normal_samples + expert_samples